
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import PolynomialFeatures
from sklearn.metrics import mean_squared_error
from sklearn.model_selection import train_test_split
from sklearn.linear_model import Lasso

x = np.random.uniform(-3, 3, size=100)
X = x.reshape(-1, 1)
y = 0.5 + x ** 2 + x + 2 + np.random.normal(0, 1, size=100)
plt.scatter(x, y)
plt.show()

def PolynomialRegression(degree):
    return Pipeline([
        ('poly', PolynomialFeatures(degree=degree)),
        ('std_scaler', StandardScaler()),
        ('lin_reg', LinearRegression())
    ])


def LassoRegression(degree, alpha):
    return Pipeline([
        ('poly', PolynomialFeatures(degree=degree)),
        ('std_scaler', StandardScaler()),
        ('lasso_reg', Lasso(alpha=alpha))
    ])


np.random.seed(555)
X_train, X_test, y_train, y_test = train_test_split(X, y)

lasso_reg1 = LassoRegression(30, 0.1)
lasso_reg1.fit(X_train, y_train)
y1_predict = lasso_reg1.predict(X_test)
print(mean_squared_error(y_test, y1_predict))


poly30_reg = PolynomialRegression(degree=30)
poly30_reg.fit(X_train, y_train)
y30_predict = poly30_reg.predict(X_test)
print(mean_squared_error(y_test, y30_predict))

X_plot = np.linspace(-3, 3, 100).reshape(100, 1)
y_plot = poly30_reg.predict(X_plot)
plt.scatter(X, y)
plt.plot(X_plot[:, 0], y_plot, color='r')
plt.axis([-3, 3, 0, 10])
plt.show()

y_plot = lasso_reg1.predict(X_plot)
plt.scatter(X, y)
plt.plot(X_plot[:, 0], y_plot, color='r')
plt.axis([-3, 3, 0, 10])
plt.show()
